{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Import the math function for calculations\n",
    "import math\n",
    "import tensorflow as tf\n",
    "import numpy as np\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Class that defines the behavior of the RBM\n",
    "class RBM(object):\n",
    "    \n",
    "    def __init__(self, input_size, output_size, lr=1.0, batchsize=100):\n",
    "        \"\"\"\n",
    "        m: Number of neurons in visible layer\n",
    "        n: number of neurons in hidden layer\n",
    "        \"\"\"\n",
    "        #Defining the hyperparameters\n",
    "        self._input_size = input_size #Size of Visible\n",
    "        self._output_size = output_size #Size of outp\n",
    "        self.learning_rate = lr #The step used in gradient descent\n",
    "        self.batchsize = batchsize #The size of how much data will be used for training per sub iteration\n",
    "        \n",
    "        #Initializing weights and biases as matrices full of zeroes\n",
    "        self.w = tf.zeros([input_size, output_size], tf.float32) #Creates and initializes the weights with 0\n",
    "        self.hb = tf.zeros([output_size], tf.float32) #Creates and initializes the hidden biases with 0\n",
    "        self.vb = tf.zeros([input_size], tf.float32) #Creates and initializes the visible biases with 0\n",
    "\n",
    "\n",
    "    #Fits the result from the weighted visible layer plus the bias into a sigmoid curve\n",
    "    def prob_h_given_v(self, visible, w, hb):\n",
    "        #Sigmoid \n",
    "        return tf.nn.sigmoid(tf.matmul(visible, w) + hb)\n",
    "\n",
    "    #Fits the result from the weighted hidden layer plus the bias into a sigmoid curve\n",
    "    def prob_v_given_h(self, hidden, w, vb):\n",
    "        return tf.nn.sigmoid(tf.matmul(hidden, tf.transpose(w)) + vb)\n",
    "    \n",
    "    #Generate the sample probability\n",
    "    def sample_prob(self, probs):\n",
    "        return tf.nn.relu(tf.sign(probs - tf.random.uniform(tf.shape(probs))))\n",
    "\n",
    "    #Training method for the model\n",
    "    def train(self, X, epochs=5):\n",
    "               \n",
    "        loss = []\n",
    "        for epoch in range(epochs):\n",
    "            #For each step/batch\n",
    "            for start, end in zip(range(0, len(X), self.batchsize),range(self.batchsize,len(X), self.batchsize)):\n",
    "                batch = X[start:end]\n",
    "                    \n",
    "                #Initialize with sample probabilities\n",
    "                    \n",
    "                h0 = self.sample_prob(self.prob_h_given_v(batch, self.w, self.hb))\n",
    "                v1 = self.sample_prob(self.prob_v_given_h(h0, self.w, self.vb))\n",
    "                h1 = self.prob_h_given_v(v1, self.w, self.hb)\n",
    "                    \n",
    "                #Create the Gradients\n",
    "                positive_grad = tf.matmul(tf.transpose(batch), h0)\n",
    "                negative_grad = tf.matmul(tf.transpose(v1), h1)\n",
    "                    \n",
    "                #Update learning rates \n",
    "                self.w = self.w + self.learning_rate *(positive_grad - negative_grad) / tf.dtypes.cast(tf.shape(batch)[0],tf.float32)\n",
    "                self.vb = self.vb +  self.learning_rate * tf.reduce_mean(batch - v1, 0)\n",
    "                self.hb = self.hb +  self.learning_rate * tf.reduce_mean(h0 - h1, 0)\n",
    "                    \n",
    "            #Find the error rate\n",
    "            err = tf.reduce_mean(tf.square(batch - v1))\n",
    "            print ('Epoch: %d' % epoch,'reconstruction error: %f' % err)\n",
    "            loss.append(err)\n",
    "                    \n",
    "        return loss\n",
    "        \n",
    "    #Create expected output for our DBN\n",
    "    def rbm_output(self, X):\n",
    "        out = tf.nn.sigmoid(tf.matmul(X, self.w) + self.hb)\n",
    "        return out\n",
    "    \n",
    "    def rbm_reconstruct(self,X):\n",
    "        h = tf.nn.sigmoid(tf.matmul(X, self.w) + self.hb)\n",
    "        reconstruct = tf.nn.sigmoid(tf.matmul(h, tf.transpose(self.w)) + self.vb)\n",
    "        return reconstruct\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "RESHAPED = 784\n",
    "NB_CLASSES = 10   # number of outputs = number of digits\n",
    "\n",
    "(train_data, Y_train), (test_data, Y_test) =  tf.keras.datasets.mnist.load_data()\n",
    "train_data = train_data/np.float32(255)\n",
    "train_data = np.reshape(train_data, (train_data.shape[0], RESHAPED))\n",
    "Y_train = tf.keras.utils.to_categorical(Y_train, NB_CLASSES)\n",
    "Y_test = tf.keras.utils.to_categorical(Y_test, NB_CLASSES)\n",
    "\n",
    "test_data = test_data/np.float32(255)\n",
    "test_data = np.reshape(test_data, (test_data.shape[0], RESHAPED))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RBM:  0   784 -> 500\n",
      "RBM:  1   500 -> 200\n",
      "RBM:  2   200 -> 50\n"
     ]
    }
   ],
   "source": [
    "RBM_hidden_sizes = [500, 200 , 50 ] #create 2 layers of RBM with size 400 and 100\n",
    "\n",
    "#Since we are training, set input as training data\n",
    "inpX = train_data\n",
    "\n",
    "#Create list to hold our RBMs\n",
    "rbm_list = []\n",
    "\n",
    "#Size of inputs is the number of inputs in the training set\n",
    "input_size = train_data.shape[1]\n",
    "\n",
    "#For each RBM we want to generate\n",
    "for i, size in enumerate(RBM_hidden_sizes):\n",
    "    print ('RBM: ',i,' ',input_size,'->', size)\n",
    "    rbm_list.append(RBM(input_size, size))\n",
    "    input_size = size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "New RBM:\n",
      "Epoch: 0 reconstruction error: 0.057573\n",
      "Epoch: 1 reconstruction error: 0.051418\n",
      "Epoch: 2 reconstruction error: 0.048383\n",
      "Epoch: 3 reconstruction error: 0.046159\n",
      "Epoch: 4 reconstruction error: 0.045744\n",
      "New RBM:\n",
      "Epoch: 0 reconstruction error: 0.028117\n",
      "Epoch: 1 reconstruction error: 0.025024\n",
      "Epoch: 2 reconstruction error: 0.022722\n",
      "Epoch: 3 reconstruction error: 0.022334\n",
      "Epoch: 4 reconstruction error: 0.021838\n",
      "New RBM:\n",
      "Epoch: 0 reconstruction error: 0.053042\n",
      "Epoch: 1 reconstruction error: 0.050178\n",
      "Epoch: 2 reconstruction error: 0.049923\n",
      "Epoch: 3 reconstruction error: 0.050089\n",
      "Epoch: 4 reconstruction error: 0.047836\n"
     ]
    }
   ],
   "source": [
    "#For each RBM in our list\n",
    "for rbm in rbm_list:\n",
    "    print ('New RBM:')\n",
    "    #Train a new one\n",
    "    rbm.train(tf.cast(inpX,tf.float32)) \n",
    "    #Return the output layer\n",
    "    inpX = rbm.rbm_output(inpX)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:new_tf]",
   "language": "python",
   "name": "conda-env-new_tf-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
